#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import pandas as pd
import numpy as np

from statsmodels.regression.linear_model import OLS
import statsmodels.api as sm

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

db = 'gp_factors'
save_db = 'gp_factors/barra_stripping'

data_reader = DataReader(db=db)
save_data_reader = DataReader(db=save_db)


barra_data_dir = './zbc_factors_lib/jq_barra_data'

class FactorBarraStripping(object):
    def __init__(self):
        self.jq_barra_p1, \
        self.jq_barra_p2 = self._get_barra_data()

        self.p1_barra_columns = self.jq_barra_p1.columns
        self.p2_barra_columns = self.jq_barra_p2.columns

        self.p1_end_date = '2019-11-30'

    def _get_barra_data(self):
        new_processed_jq_barra_data_v0 = pd.read_hdf(os.path.join(barra_data_dir, 'new_processed_jq_barra_data_v0.h5'))
        new_processed_jq_barra_data_v1 = pd.read_hdf(os.path.join(barra_data_dir, 'new_processed_jq_barra_data_v1.h5'))

        new_processed_jq_barra_data_v0 = new_processed_jq_barra_data_v0.set_index(['date', 'symbol'])
        new_processed_jq_barra_data_v1 = new_processed_jq_barra_data_v1.set_index(['date', 'symbol'])

        return new_processed_jq_barra_data_v0, new_processed_jq_barra_data_v1

    def factor_barra_stripping(self, factor_name, is_update=False, is_overwrite=False):
        ## TODO - 分段回归
        # 读取因子数据
        # factor_name = 'zg02_ths_user_hehavior_factor_mining_v2_v30'
        factor_data = data_reader.read_factor_table(factor_name)

        # p1_end_date = '2019-11-30'
        factor_data_p1 = factor_data[factor_data['date'] <= self.p1_end_date]
        factor_data_p2 = factor_data[factor_data['date'] > self.p1_end_date]

        factor_data_p1 = factor_data_p1.set_index(['date', 'stock_code'])
        factor_data_p2 = factor_data_p2.set_index(['date', 'stock_code'])

        # P1
        factor_data_p1 = pd.concat([factor_data_p1, self.jq_barra_p1], axis=1, join_axes=[factor_data_p1.index])
        factor_data_p1 = factor_data_p1.dropna()

        model = OLS(factor_data_p1[factor_name], factor_data_p1[self.p1_barra_columns], hasconst=True)

        res = model.fit(params_only=True)

        # P1剥离
        factor_data_p1[factor_name] = \
            factor_data_p1[factor_name] - factor_data_p1[self.p1_barra_columns].values @ res.params.values

        # P2
        factor_data_p2 = pd.concat([factor_data_p2, self.jq_barra_p2], axis=1, join_axes=[factor_data_p2.index])
        factor_data_p2 = factor_data_p2.dropna()

        model = OLS(factor_data_p2[factor_name], factor_data_p2[self.p2_barra_columns], hasconst=True)

        res = model.fit(params_only=True)

        # P2剥离
        factor_data_p2[factor_name] = \
            factor_data_p2[factor_name] - factor_data_p2[self.p2_barra_columns].values @ res.params.values

        # concat
        factor_data = pd.concat([factor_data_p1[[factor_name]],
                                 factor_data_p2[[factor_name]]], axis=0)

        factor_data = factor_data.reset_index()

        del factor_data_p1
        del factor_data_p2

        # save
        save_data_reader.update_data(factor_data)

        print(factor_name, 'barra stripped!\n')

        if not is_update:
            # save_data_reader.update_data(factor_data[factor_data['date'] <= '2020-03-01'])
            save_data_reader.create_factor_table(factor_name, main_columns=['date', 'stock_code'])
            print(factor_name, 'barra stripped created!\n')
        else:
            # save_data_reader.update_data(factor_data[factor_data['date'] > '2020-03-01'])
            save_data_reader.update_factor_data_table(factor_name,
                                                      main_columns=['date', 'stock_code'],
                                                      filter_column='date',
                                                      overwrite=is_overwrite,
                                                      drop_duplicates_by=['stock_code', 'date'],
                                                      is_filter_duplicates=True)

            print(factor_name, 'barra stripped updated!\n')

    def multi_factors_barra_stripping(self, factor_name_list, is_update=False, is_overwrite=False):
        for factor_name in factor_name_list:
            self.factor_barra_stripping(factor_name, is_update=is_update, is_overwrite=is_overwrite)


if __name__ == '__main__':
    is_update = False
    is_overwrite = False

    factor_barra_strip_api = FactorBarraStripping()

    ##
    factor_name_list = [
        # 'zbc_ths_user_hehavior_factor_mining_v2_v31'
        # 'zbc_ths_user_hehavior_factor_mining_v2_v%d' % id for id in range(1, 31)
        'zbc_ths_user_hehavior_factor_mining_v2_v%d' % id for id in range(31, 42)
    ]

    factor_barra_strip_api.multi_factors_barra_stripping(factor_name_list=factor_name_list,
                                                         is_update=is_update,
                                                         is_overwrite=is_overwrite)


